import pandas as pd
import matplotlib.pyplot as plt

# 定义填补函数
def fill_sediment(row, df):
    if pd.isna(row['含沙量(kg/m3) ']):
        # 获取当前行的位置
        index = df.index.tolist().index(row.name)
        d_forward = 1
        d_backward = 1
        # 获取 DataFrame 的最大位置
        max_index = len(df) - 1
        min_index = 0

        # 查找下一个非空值
        next_non_na_index = None
        while index + d_forward <= max_index:
            if not pd.isna(df.iloc[index + d_forward]['含沙量(kg/m3) ']):
                next_non_na_index = index + d_forward
                break
            d_forward += 1

        # 查找前一个非空值
        prev_non_na_index = None
        while index - d_backward >= min_index:
            if not pd.isna(df.iloc[index - d_backward]['含沙量(kg/m3) ']):
                prev_non_na_index = index - d_backward
                break
            d_backward += 1

        # 如果找到了前后非空值
        if next_non_na_index is not None and prev_non_na_index is not None:
            q_st_prev = df.iloc[prev_non_na_index]['含沙量(kg/m3) ']
            q_st_next = df.iloc[next_non_na_index]['含沙量(kg/m3) ']
            return (d_forward / (d_forward + d_backward)) * q_st_prev + (d_backward / (d_forward + d_backward)) * q_st_next
        # 如果只找到了前一个非空值
        elif prev_non_na_index is not None:
            return df.iloc[prev_non_na_index]['含沙量(kg/m3) ']
        # 如果只找到了下一个非空值
        elif next_non_na_index is not None:
            return df.iloc[next_non_na_index]['含沙量(kg/m3) ']

    return row['含沙量(kg/m3) ']

def data_clean():
    target_years = [2017, 2018, 2019, 2021]
    # 读取 Excel 文件
    file_path = 'ori_message.xlsx'
    try:
        df_all = pd.read_excel(file_path, sheet_name=None)
    except FileNotFoundError:
        print(f"错误：文件 {file_path} 未找到。")
    except Exception as e:
        print(f"错误：读取文件时发生未知错误：{e}")
    else:
        all_6_year_data = []
        for year in df_all:
            print(year)
            df = df_all[year]

            # 填充年、月、日列
            df[['年', '月', '日']] = df[['年', '月', '日']].ffill()

            # 筛选每天 8:00 的数据
            df_8am = df[df['时间'] == '8:00']
            # 将年、月、日列转换为整数类型
            df_8am[['年', '月', '日']] = df_8am[['年', '月', '日']].astype(int)
            df_8am.drop(['时间'], axis=1, inplace=True)

            all_dates = []
            for month in range(1, 13):
                last_day = pd.Timestamp(f'{year}-{month}-01').days_in_month
                for day in range(1, last_day + 1):
                    if not (month == 2 and day == 29):
                        all_dates.append((year, month, day))
            all_dates_df = pd.DataFrame(all_dates, columns=['年', '月', '日'])
            # 确保 all_dates_df 中的年、月、日列也是整数类型
            all_dates_df[['年', '月', '日']] = all_dates_df[['年', '月', '日']].astype(int)
            # 合并数据，缺失数据用空值填充
            df_merged = pd.merge(all_dates_df, df_8am, on=['年', '月', '日'], how='left')
            df_merged = df_merged[~((df_merged['月'] == 2) & (df_merged['日'] == 29))]
            # 填补含沙量缺失值
            df_merged['含沙量(kg/m3) '] = df_merged.apply(lambda d: fill_sediment(d, df_merged), axis=1)
            # 去除含沙量列存在缺失值的行
            df_merged = df_merged.dropna(subset=['含沙量(kg/m3) '])
            # 设置显示选项以打印所有行
            pd.set_option('display.max_rows', None)
            print(df_merged)
            all_6_year_data.append(df_merged)

        # 创建新的 DataFrame 来存储最终结果
        new_table = {}
        for df in all_6_year_data:
            year = df['年'].iloc[0]
            # 提取含沙量列的值
            sediment_values = df['含沙量(kg/m3) '].tolist()
            new_table[year] = sediment_values

        # 将字典转换为 DataFrame
        result_df = pd.DataFrame(new_table).T
        # 设置列名，以 'Day_1', 'Day_2', ... 命名
        result_df.columns = [f'Day_{i + 1}' for i in range(result_df.shape[1])]

        # 计算 6 年的每日平均含沙量
        average_sediment = result_df.sum(axis=0) / len(result_df)

        print("6 年的每日平均含沙量：")
        print(average_sediment)

        # 生成日期标签列表（不包含 2 月 29 号）
        date_labels = []
        for month in range(1, 13):
            last_day = pd.Timestamp(f'2023-{month}-01').days_in_month  # 用非闰年的年份来获取每月天数
            for day in range(1, last_day + 1):
                if not (month == 2 and day == 29):
                    date_str = f"{month:02d}-{day:02d}"
                    date_labels.append(date_str)

        # 绘制折线图
        plt.figure(figsize=(15, 6))
        plt.plot(date_labels, average_sediment.values, marker='o', linestyle='-')

        # 设置图表标题和坐标轴标签
        plt.title('6 年的每日平均含沙量')
        plt.xlabel('日期（月 - 日）')
        plt.ylabel('含沙量 (kg/m³)')
        plt.xticks(rotation=90)
        plt.grid(True)

        # 显示图表
        plt.tight_layout()
        plt.show()

data_clean()